# 构建一个关于《玫瑰的故事》知识库问答系统
# 系统只能回答关于玫瑰的故事的内容，其他内容一概不进行回复
# 需要给系统喂数据，让大模型只能根据我们喂的数据进行只能回复
# 知识库其实说白了就是一个txt文本文件
# 喂数据就是把数据喂给了langchain中prompt
import base64
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
import streamlit as st
import model.model as mm
import data.data as dd
import datetime
st.set_page_config(
    page_title="知识库工作页面",
    page_icon="🤖"
)
mm.page_config()
st.html("""
        <style>
        button{
            width:100% !important;
        }
        #root > div:nth-child(1) > div.withScreencast > div > div > div > section.st-emotion-cache-1gv3huu.eczjsme18 > div.st-emotion-cache-6qob1r.eczjsme11 > div.st-emotion-cache-1gwvy71.eczjsme12 > div > div > div > div > div:nth-child(1) > div > div > h1{
        text-align:center;
        }
    </style>
    """)
llm = ChatOpenAI(
    model="glm-4-0520",
    temperature=0.2,
    api_key="f329dd27aec19d68aff2fe9feab29e6a.FmBng9a9FWUcG8gM",
    base_url="https://open.bigmodel.cn/api/paas/v4"
)
prompt = PromptTemplate(
    input_variables=["content","problem"],
    template="""
        你是一个专业知识库问答专家，你需要从提供的知识库{content}中检索答案，如果知识库中没有这个内容，不要自己发挥，也不要自己推理，只能根据知识库已有的内容来回答，请不需要自己发挥{problem}
    """
)

chain = prompt | llm

user_id = st.session_state.user_id
username = st.session_state.username
# 构建一个通用性的知识库问答系统，需要用户自己上传他想提问的数据库
file = st.file_uploader("请选择你的知识库文件",type=["txt"])
if file is not None:
    content = file.read().decode("utf-8")
    input = st.chat_input("请输入你的问题")
    if input:
        with st.chat_message("user"):
            st.write(input)
            dd.insert_work_message(user_id,input,"user",datetime.datetime.now())
        result = chain.invoke({
            "content": content,
            "problem": input
        })
        with st.chat_message("assistant"):
            st.write(result.content)
            res = result.content
            dd.insert_work_message(user_id,res,"assistant",datetime.datetime.now())
def main_bg(main_bg):
    main_bg_ext = "png"
    st.markdown(
        f"""
         <style>
         .stApp {{
             background: url(data:image/{main_bg_ext};base64,{base64.b64encode(open(main_bg, "rb").read()).decode()});
             background-size: cover
         }}
         </style>
         """,
        unsafe_allow_html=True
    )

# 调用
main_bg('images/login6.png')



